#include <stdio.h>
#include <stdlib.h>
#include <signal.h>
#include <unistd.h>
#include "mytbf.h"

typedef void (*sighandler_t)(int);

static struct mytbf_st *job[MYTBF_MAX];
static int inited = 0;
static sighandler_t alrm_handler_save;

struct mytbf_st
{
  int cps; // number of characters per second
  int burst; // max number of token
  int token; // number of token to send characters
  int pos;
};

static int get_free_pos(void)
{
  for (int i = 0; i < MYTBF_MAX; i++) {
    if (job[i] == NULL) {
      return i;
    }
  }
  return -1;
}

static void alrm_handler(int s)
{
  alarm(1);
  for (int i = 0; i < MYTBF_MAX; i++) {
    if(job[i] != NULL)
    {
      job[i]->token += job[i]->cps;
      if(job[i]->token > job[i]->burst) {
        job[i]->token = job[i]->burst;
      }
    }
  }
}

static void module_unload(void)
{
  signal(SIGALRM, alrm_handler_save);
  // close the registered alarm
  alarm(0);

  for (int i = 0; i < MYTBF_MAX; i++) {
    free(job[i]);
  }
}

static void module_load(void)
{
  alrm_handler_save = signal(SIGALRM, alrm_handler);
  alarm(1);

  // register a function to be called at normal process termination
  atexit(module_unload);
}

mytbf_t *mytbf_init(int cps, int burst)
{
  int pos;
  pos = get_free_pos();
  if (pos < 0) {
    return NULL;
  }

  if (!inited) {
    module_load();
    inited = 1;
  }

  struct mytbf_st *me = (struct mytbf_st *)malloc(sizeof(struct mytbf_st));
  if (me == NULL) {
    return NULL;
  }
  me->token = 0;
  me->cps = cps;
  me->burst = burst;
  me->pos = pos;
  job[pos] = me;

  return me;
}

static int min(int a, int b)
{
  return (a > b) ? b : a;
}

int mytbf_fetchtoken(mytbf_t *tbf, int nums)
{
  int n;
  struct mytbf_st *me = (struct mytbf_st *)tbf;
  if (nums <= 0) {
    return -1;
  }

  while (me->token <= 0) {
    // wait for signal
    pause();
  }

  n = min(me->token, nums);
  me->token -= n;
  return n;
}

int mytbf_returntoken(mytbf_t *tbf, int nums)
{
  struct mytbf_st *me = (struct mytbf_st *)tbf;
  if (nums <= 0) {
    return -1;
  }

  me->token += nums;
  if (me->token > me->burst) {
    me->token = me->burst;
  }

  return nums;
}

int mytbf_destroy(mytbf_t *tbf)
{
  struct mytbf_st *me = (struct mytbf_st *)tbf;
  job[me->pos] = NULL;
  free(tbf);
  return 0;
}
